package main

import (
	"context"
	"flag"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	redisAddrFlag     = flag.String("redis.addr", "", "Redis server address (env REDIS_ADDR)")
	redisPasswordFlag = flag.String("redis.password", "", "Redis password (env REDIS_PASSWORD)")
	redisKeyFlag      = flag.String("redis.key", "", "Redis key(s) to query, comma separated (env REDIS_KEY)")
	listenAddrFlag    = flag.String("web.listen-address", "", "Address to listen on for web interface and telemetry (env LISTEN_ADDR)")
)

var redisKeyGauge = prometheus.NewGaugeVec(prometheus.GaugeOpts{
	Name: "redis_key_value",
	Help: "Value of the specified Redis key",
}, []string{"redis_addr", "key"})

// 支持多个 key

type redisCollector struct {
	client    *redis.Client
	keys      []string
	redisAddr string
}

func (r *redisCollector) Describe(ch chan<- *prometheus.Desc) {
	redisKeyGauge.Describe(ch)
}

func (r *redisCollector) Collect(ch chan<- prometheus.Metric) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	for _, key := range r.keys {
		val, err := r.client.Get(ctx, key).Result()
		labels := prometheus.Labels{"redis_addr": r.redisAddr, "key": key}
		if err == redis.Nil {
			redisKeyGauge.With(labels).Set(0)
		} else if err != nil {
			log.Printf("Error getting key %s from Redis: %v", key, err)
			redisKeyGauge.With(labels).Set(0)
		} else {
			f, err := strconv.ParseFloat(val, 64)
			if err != nil {
				log.Printf("Value of key %s is not a float: %v", key, err)
				redisKeyGauge.With(labels).Set(0)
			} else {
				redisKeyGauge.With(labels).Set(f)
			}
		}
	}
	redisKeyGauge.Collect(ch)
}

func getEnvOrFallback(flagVal, envKey, defaultVal string) string {
	if flagVal != "" {
		return flagVal
	}
	envVal := os.Getenv(envKey)
	if envVal != "" {
		return envVal
	}
	return defaultVal
}

func main() {
	flag.Parse()

	redisAddr := getEnvOrFallback(*redisAddrFlag, "REDIS_ADDR", "127.0.0.1:6379")
	redisPassword := getEnvOrFallback(*redisPasswordFlag, "REDIS_PASSWORD", "")
	redisKeysStr := getEnvOrFallback(*redisKeyFlag, "REDIS_KEY", "")
	listenAddr := getEnvOrFallback(*listenAddrFlag, "LISTEN_ADDR", ":9100")

	if redisKeysStr == "" {
		log.Fatalf("redis.key must be specified either by flag or environment variable REDIS_KEY")
	}

	// 支持多个 key，逗号分隔
	var redisKeys []string
	for _, k := range splitAndTrim(redisKeysStr, ",") {
		if k != "" {
			redisKeys = append(redisKeys, k)
		}
	}
	if len(redisKeys) == 0 {
		log.Fatalf("No valid redis keys specified")
	}

	rdb := redis.NewClient(&redis.Options{
		Addr:     redisAddr,
		Password: redisPassword,
	})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := rdb.Ping(ctx).Err(); err != nil {
		log.Fatalf("Failed to connect to Redis at %s: %v", redisAddr, err)
	}

	collector := &redisCollector{
		client:    rdb,
		keys:      redisKeys,
		redisAddr: redisAddr,
	}
	prometheus.MustRegister(collector)

	http.Handle("/metrics", promhttp.Handler())

	log.Printf("Starting server at %s, querying Redis %s keys %v", listenAddr, redisAddr, redisKeys)
	if err := http.ListenAndServe(listenAddr, nil); err != nil {
		log.Fatalf("HTTP server error: %v", err)
	}
}

// splitAndTrim 用于分割字符串并去除空格
func splitAndTrim(s, sep string) []string {
	var result []string
	for _, part := range split(s, sep) {
		trimmed := trimSpace(part)
		result = append(result, trimmed)
	}
	return result
}

func split(s, sep string) []string {
	var res []string
	start := 0
	for i := 0; i+len(sep) <= len(s); {
		if s[i:i+len(sep)] == sep {
			res = append(res, s[start:i])
			start = i + len(sep)
			i = start
		} else {
			i++
		}
	}
	res = append(res, s[start:])
	return res
}

func trimSpace(s string) string {
	start := 0
	end := len(s)
	for start < end && (s[start] == ' ' || s[start] == '\t' || s[start] == '\n' || s[start] == '\r') {
		start++
	}
	for end > start && (s[end-1] == ' ' || s[end-1] == '\t' || s[end-1] == '\n' || s[end-1] == '\r') {
		end--
	}
	return s[start:end]
}
